#include "ublk.h"
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <fcntl.h>
#include <unistd.h>
#include <stdlib.h>
#include <errno.h>
#include <liburing/io_uring.h>

// ***** 自包含 ublk 和 io_uring 的现代 API *****
#define UBLK_CMD_ADD_DEV        _IOW('u', 0, struct ublk_params)
#define UBLK_CMD_DEL_DEV        _IOW('u', 1, __u32)
#define UBLK_CMD_START_DEV      _IOW('u', 2, __u32)
#define UBLK_CMD_GET_DEV_FD     _IOWR('u', 4, __u32)

// 这个结构体在您的 liburing 头文件中缺失
struct io_uring_cmd {
        __u32   cmd_op;
        __u32   flags; // IORING_URING_CMD_FIXED
        __u64   __pad0;
        __u64   addr;
        __u32   len;
        __u32   __pad1;
        __u64   user_data;
        __u16   tag; // internal use
        __u16   __pad2;
        __u64   __pad3[2];
};

struct ublk_params {
    __u32 dev_id;
    __u32 nr_hw_queues;
    __u32 queue_depth;
    __u64 reserved[7];
};
// ***** API 定义结束 *****


static void complete_io(io_request_t *req, bool success);
static void handle_io_request(void* arg);
// ... (所有辅助函数保持不变)
static void process_disk_read(void* arg);
static void on_chunk_fetched(void* arg, bool success);

static void complete_io(io_request_t *req, bool success) {
    struct io_uring_sqe *sqe = io_uring_get_sqe(&req->daemon->ring);
    struct io_uring_cmd cmd = {
        .cmd_op = UBLK_IO_COMMIT_AND_FETCH_REQ,
        .tag = req->tag,
        .user_data = success ? 0 : -EIO,
    };
    io_uring_prep_rw(IORING_OP_URING_CMD, sqe, req->daemon->ublk_fd, &cmd, sizeof(cmd), 0);
    io_uring_sqe_set_data(sqe, NULL);
    io_uring_submit(&req->daemon->ring);
    free(req);
}

static void on_chunk_fetched(void* arg, bool success) {
    fetch_task_arg_t *task_arg = (fetch_task_arg_t *)arg;
    io_request_t *req = (io_request_t*)task_arg->callback_arg;
    if (success) {
        cache_manager_write_to_disk(req->daemon->cm, task_arg->hash, task_arg->buffer, CHUNK_SIZE);
        memcpy(req->user_ptr, task_arg->buffer, CHUNK_SIZE);
    }
    complete_io(req, success);
    free(task_arg->buffer); free(task_arg->hash); free(task_arg);
}

static void process_disk_read(void* arg) {
    io_request_t *req = (io_request_t *)arg;
    const char* hash = cache_manager_get_hash(req->daemon->cm, req->lba);
    if (cache_manager_read_from_disk(req->daemon->cm, hash, req->user_ptr)) {
        complete_io(req, true);
    } else {
        fetcher_fetch_chunk(req->daemon->fetcher, hash, on_chunk_fetched, req);
    }
}

static void handle_io_request(void* arg) {
    io_request_t *req = (io_request_t *)arg;
    const char* hash = cache_manager_get_hash(req->daemon->cm, req->lba);
    if (!hash) {
        memset(req->user_ptr, 0, CHUNK_SIZE);
        complete_io(req, true);
        return;
    }
    thread_pool_add(req->daemon->worker_pool, process_disk_read, req);
}


int ublk_daemon_run(int dev_id, const char* lba_map_path, const char* cache_dir, const char* registry, const char* repo) {
    ublk_daemon_t daemon;
    daemon.dev_id = dev_id;
    daemon.cm = cache_manager_create(cache_dir, lba_map_path);
    daemon.fetcher = fetcher_create(registry, repo, 4);
    daemon.worker_pool = thread_pool_create(8, 256);
    if (!daemon.cm || !daemon.fetcher || !daemon.worker_pool) { LOG_ERROR("守护进程初始化失败"); return -1; }

    int ctrl_fd = open("/dev/ublk-control", O_RDWR);
    if (ctrl_fd < 0) { LOG_ERROR("无法打开 /dev/ublk-control"); return -1; }
    
    struct ublk_params p = {
        .dev_id = dev_id,
        .nr_hw_queues = 1,
        .queue_depth = 128,
    };
    if (ioctl(ctrl_fd, UBLK_CMD_ADD_DEV, &p) < 0) {
        LOG_ERROR("UBLK_CMD_ADD_DEV 失败");
        close(ctrl_fd);
        return -1;
    }
    
    daemon.ublk_fd = ioctl(ctrl_fd, UBLK_CMD_GET_DEV_FD, &dev_id);
    if (daemon.ublk_fd < 0) {
        LOG_ERROR("UBLK_CMD_GET_DEV_FD 失败");
        ioctl(ctrl_fd, UBLK_CMD_DEL_DEV, &dev_id);
        close(ctrl_fd);
        return -1;
    }

    io_uring_queue_init(p.queue_depth, &daemon.ring, 0);
    
    // 从 ubk_fd 读取启动命令
    struct io_uring_cmd start_cmd;
    read(daemon.ublk_fd, &start_cmd, sizeof(start_cmd));
    if (start_cmd.cmd_op != UBLK_CMD_OP_START_DEV) { LOG_ERROR("启动 ublk 设备失败"); return -1; }

    LOG_INFO("ublk 设备 %d 启动成功, ublk_fd: %d. 开始 I/O 循环...", dev_id, daemon.ublk_fd);
    
    while (1) {
        struct io_uring_sqe *sqe = io_uring_get_sqe(&daemon.ring);
        struct io_uring_cmd fetch_cmd = { .cmd_op = UBLK_IO_FETCH_REQ };
        io_uring_prep_rw(IORING_OP_URING_CMD, sqe, daemon.ublk_fd, &fetch_cmd, sizeof(fetch_cmd), 0);
        io_uring_sqe_set_data(sqe, NULL);
        io_uring_submit(&daemon.ring);

        struct io_uring_cqe *cqe;
        io_uring_wait_cqe(&daemon.ring, &cqe);
        
        if (cqe->res >= 0) {
            const struct ublksrv_io_data *iod = (const struct ublksrv_io_data*)cqe->user_data;
            const struct ublksrv_io_desc *desc = (const struct ublksrv_io_desc *)iod->addr;
            
            if (ublksrv_get_op(desc) == UBLK_IO_OP_READ) {
                io_request_t *req = (io_request_t*)malloc(sizeof(io_request_t));
                req->daemon = &daemon;
                req->tag = iod->tag;
                req->lba = desc->start_sector / (CHUNK_SIZE / 512);
                req->len = desc->nr_sectors * 512;
                req->user_ptr = (void*)(unsigned long)iod->addr;
                handle_io_request(req);
            }
        }
        io_uring_cqe_seen(&daemon.ring, cqe);
    }
    
    // 清理
    return 0;
}